#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import device_manager.arm_manager 
import device_manager.gripper_manager 
import numpy as np
import rospy
import json
import os
import cv2
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError
from rospy_message_converter import message_converter

class GetData:
    def __init__(self):
        self.save_data_flag = False
        self.arm  = device_manager.arm_manager.ArmManager()
        self.arm.set_max_acceleration_and_velocity()
        self.__bridge = CvBridge()
        self.file_hdir = os.path.join(os.path.expanduser('~'), "data")
        self.orientation = [0,0,0,0]
        self.count = 0
        self.pcount = [0,0,0]
        self.raw_position = [0,0,0]

    def moving_arm(self):

        position = self.raw_position
        if self.pcount[1] <3:
            position[0] = 0.3*pcount[0] + self.raw_position[0]
            position[1] = 0.3*pcount[1] + self.raw_position[1]
            self.pcount[1] = pcount[1]+1
        else:
            self.pcount[1]=0
            self.pcount[0]=pcount[0]+1
            position[0] = 0.3*pcount[0] + self.raw_position[0]
            position[1] = 0.3*pcount[1] + self.raw_position[1]
            self.pcount[1]= self.pcount[1]+1
        self.arm.set_target_pose(position,self.orientation)
        rospy.sleep(1)
        self.save_data_flag = True


    def save_data(self, bgr):
       
        pose = self.arm.get_current_pose()

        extend1 = (str(self.count)+".jpg")
        extend2 = (str(self.count)+".json")

        im_name = os.path.join(self.file_hdir, extend1)
        jp_name = os.path.join(self.file_hdir, extend2)

        pose_dict = message_converter.convert_ros_message_to_dictionary(pose)
        fp = open(jp_name, 'w+')
        fp.write(json.dumps(pose_dict)+'\n')
        fp.close()

        cv2.imwrite(im_name, bgr)

        self.count = self.count+1
        self.save_data_flag = False
        rospy.sleep(0.2)

    def __image_callback(self, msg_image):
        
        bgr = self.__bridge.imgmsg_to_cv2(msg_image, 'bgr8')
        if self.count<9:
            if self.save_data_flag:
                self.save_data(bgr)
            else:
                self.moving_arm()
        if self.count==9:
            rospy.INFO("[INFO]   data has been saved.")
            self.count = self.count+1

    def get_calibration_data(self):

        # self.arm.back_to_zero()
        
        pose = self.arm.get_current_pose()
        self.orientation[0]= pose.orientation.x
        self.orientation[1]= pose.orientation.y
        self.orientation[2]= pose.orientation.z
        self.orientation[3]= pose.orientation.w
        
        self.raw_position[0] = pose.position.x
        self.raw_position[1] = pose.position.y
        self.raw_position[2] = pose.position.z

        image_raw_topic_str = rospy.get_param('~image_src','/image_raw') 
        rospy.Subscriber(image_raw_topic_str, Image, self.__image_callback, queue_size=1)
        
        bgr = cv2.imread("/home/ly/Downloads/1.jpg")
        self.save_data(bgr)
        rospy.spin()

if __name__ == '__main__':
    
    rospy.init_node('DEVICE_MANAGER', anonymous=False)
    gcd = GetData()
    gcd.get_calibration_data()

    
   
